import cv2  # 导入cv2模块,用于读取和处理图像
import numpy as np  # 导入numpy模块,主要用于数值计算
import os  # 导入os模块,用于处理文件和目录
from os.path import exists  # 从os.path模块导入exists函数,用于检测文件或目录是否存在
from imutils import paths  # 导入imutils中的paths工具,用于获取文件路径
import pickle  # 导入pickle模块,用于序列化和反序列化Python对象结构
import zipfile
import tempfile
import shutil
import glob
from tqdm import tqdm  # 导入tqdm模块,用于在循环中添加进度条

import logging  # 用于记录日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') # 配置logging

# 定义函数get_size,获取文件大小,返回文件大小的兆字节的浮点数表示
def get_size(file):
    """
    获取指定文件的大小（以MB为单位）
    
    参数:
    file (str): 文件的路径
    
    返回:
    float: 文件大小（MB）
    """
    # 使用 os.path.getsize 获取字节数并换算为 MB
    size_bytes = os.path.getsize(file)
    return size_bytes / (1024 * 1024)

# 定义函数createXY,用于从图像创建特征(X)和标签(y)
def createXY(train_folder, dest_folder, method='vgg', batch_size=64):
    # 支持传入 zip 文件或自动在当前目录查找 zip 并解压以生成图像数据集
    extracted_tmp = None
    try:
        # 如果传入的是一个 .zip 文件，优先解压使用
        if os.path.isfile(train_folder) and train_folder.lower().endswith('.zip'):
            logging.info(f"传入的 train_folder 是 zip 文件，解压到临时目录后使用: {train_folder}")
            extracted_tmp = tempfile.mkdtemp(prefix='createXY_')
            try:
                with zipfile.ZipFile(train_folder, 'r') as zf:
                    zf.extractall(extracted_tmp)
                train_folder = extracted_tmp
            except Exception as e:
                logging.exception(f"解压传入 zip 失败: {e}")

        # 如果用户传入的路径不存在或不存在图片，尝试在当前工作目录查找 zip 文件并解压试用
        if not os.path.exists(train_folder):
            # 先尝试解析相对路径保留原逻辑
            alt = os.path.abspath(os.path.join(os.getcwd(), train_folder))
            if os.path.exists(alt):
                logging.info(f"train_folder 未直接找到，使用解析后路径: {alt}")
                train_folder = alt
            else:
                # 查找当前目录下的 zip 文件，尝试解压并检查是否包含图片
                zips = glob.glob('*.zip')
                for z in zips:
                    tmp = tempfile.mkdtemp(prefix='createXY_')
                    try:
                        with zipfile.ZipFile(z, 'r') as zf:
                            zf.extractall(tmp)
                        # 如果解压后找到图片则使用该临时目录
                        if any(paths.list_images(tmp)):
                            logging.info(f"在当前目录发现 zip 并解压可用数据: {z} -> {tmp}")
                            extracted_tmp = tmp
                            train_folder = tmp
                            break
                        else:
                            shutil.rmtree(tmp)
                    except Exception:
                        shutil.rmtree(tmp)
                        continue
    # 方法相关的文件名，避免不同特征方法覆盖同一份缓存
    x_file_path = os.path.join(dest_folder, f"X_{method}.pkl")
    y_file_path = os.path.join(dest_folder, f"y_{method}.pkl")
    # 兼容旧版缓存文件名（X.pkl / y.pkl）
        def _cleanup():
            try:
                if extracted_tmp and os.path.exists(extracted_tmp):
                    shutil.rmtree(extracted_tmp)
                    logging.info(f"清理临时解压目录: {extracted_tmp}")
            except Exception:
                pass
    legacy_x_path = os.path.join(dest_folder, "X.pkl")
    legacy_y_path = os.path.join(dest_folder, "y.pkl")

    # 如果方法专属的 X 和 y 已经存在，则直接读取；否则兼容读取旧版通用缓存
    if os.path.exists(x_file_path) and os.path.exists(y_file_path):
        x_size = os.path.getsize(x_file_path)
        y_size = os.path.getsize(y_file_path)
        logging.info(f"发现缓存（method={method}）：直接读取 {os.path.basename(x_file_path)}, {os.path.basename(y_file_path)}")
        logging.info(f"X文件大小:{get_size(x_file_path):.2f}MB")
        logging.info(f"y文件大小:{get_size(y_file_path):.2f}MB")
        # 如果任一文件为空，跳过缓存并删除空文件以避免后续误用
        if x_size == 0 or y_size == 0:
            logging.warning(f"检测到空缓存文件，忽略并删除: {x_file_path if x_size==0 else ''} {y_file_path if y_size==0 else ''}")
            try:
                if x_size == 0:
                    os.remove(x_file_path)
                if y_size == 0:
                    os.remove(y_file_path)
            except Exception:
                pass
        else:
            try:
                with open(x_file_path, 'rb') as fx:
                    X = pickle.load(fx)
                with open(y_file_path, 'rb') as fy:
                    y = pickle.load(fy)
                # 额外检查加载的数据是否确实非空
                try:
                    import numpy as _np
                    if _np.array(X).size == 0 or _np.array(y).size == 0:
                        logging.warning("缓存文件包含空数组，忽略并删除缓存文件")
                        try:
                            os.remove(x_file_path)
                            os.remove(y_file_path)
                        except Exception:
                            pass
                    else:
                        return X, y
                except Exception:
                    # 如果 numpy 检查失败，仍然尝试返回加载的数据
                    return X, y
            except Exception as e:
                logging.exception(f"读取缓存失败，忽略缓存并重新生成: {e}")
    elif os.path.exists(legacy_x_path) and os.path.exists(legacy_y_path):
        lx_size = os.path.getsize(legacy_x_path)
        ly_size = os.path.getsize(legacy_y_path)
        logging.info(f"发现旧版缓存：直接读取 {os.path.basename(legacy_x_path)}, {os.path.basename(legacy_y_path)}")
        logging.info(f"X文件大小:{get_size(legacy_x_path):.2f}MB")
        logging.info(f"y文件大小:{get_size(legacy_y_path):.2f}MB")
        if lx_size == 0 or ly_size == 0:
            logging.warning(f"检测到空旧版缓存文件，忽略并删除: {legacy_x_path if lx_size==0 else ''} {legacy_y_path if ly_size==0 else ''}")
            try:
                if lx_size == 0:
                    os.remove(legacy_x_path)
                if ly_size == 0:
                    os.remove(legacy_y_path)
            except Exception:
                pass
        else:
            try:
                with open(legacy_x_path, 'rb') as fx:
                    X = pickle.load(fx)
                with open(legacy_y_path, 'rb') as fy:
                    y = pickle.load(fy)
                try:
                    import numpy as _np
                    if _np.array(X).size == 0 or _np.array(y).size == 0:
                        logging.warning("旧版缓存包含空数组，忽略并删除缓存文件")
                        try:
                            os.remove(legacy_x_path)
                            os.remove(legacy_y_path)
                        except Exception:
                            pass
                    else:
                        return X, y
                except Exception:
                    return X, y
            except Exception as e:
                logging.exception(f"读取旧版缓存失败，忽略缓存并重新生成: {e}")

    logging.info(f"读取所有图像，生成X和y（method={method}）")  # 提示用户开始读取图像并生成X和y

    # 其余路径容错解析（已在上方处理 zip 场景）
    if not os.path.exists(train_folder):
        alt2 = os.path.abspath(os.path.join(os.getcwd(), os.path.basename(train_folder)))
        if os.path.exists(alt2):
            logging.info(f"train_folder 未直接找到，使用候选路径: {alt2}")
            train_folder = alt2

    image_paths = list(paths.list_images(train_folder))  # 获取训练文件夹中所有图像的路径

    # 如果没有找到图像，给出明确的错误信息，避免后续静默返回空数组
    if len(image_paths) == 0:
        logging.error(f"在路径 {train_folder} 中未找到任何图片文件。请检查路径或数据组织（例如 data/train/dog/*.jpg）。")
        return np.array([]), np.array([])

    X = []  # 初始化X列表,用于存储特征
    y = []  # 初始化y列表,用于存储标签
    
    # 根据传入的方法选择不同的模型（延迟导入 TensorFlow，避免 flat 模式要求安装 TF）
    if method == 'vgg':
        from tensorflow.keras.applications.vgg16 import VGG16, preprocess_input  # 延迟导入
        from tensorflow.keras.utils import load_img, img_to_array  # 延迟导入
        model = VGG16(weights='imagenet', include_top=False, pooling="max")  # 加载VGG16模型,不包括顶层,使用最大池化
        logging.info("完成构建 VGG16 模型")  # 提示用户VGG16模型构建完成
    elif method == 'flat':
        model = None  # 如果方法为'flat',不使用任何预训练模型

    # 计算需要的批次数
    num_batches = len(image_paths) // batch_size + (1 if len(image_paths) % batch_size else 0)

    # 使用进度条对批次进行循环处理
    for idx in tqdm(range(num_batches), desc="读取图像"):
        batch_images = []  # 初始化存储批次图像的列表
        batch_labels = []  # 初始化存储批次标签的列表
        
        start = idx * batch_size  # 计算批次开始的索引
        end = min((idx + 1) * batch_size, len(image_paths))  # 计算批次结束的索引

        # 对于每个批次中的图像
        for i in range(start, end):
            image_path = image_paths[i]  # 获取图像路径
            if method == 'vgg':
                # 仅在 vgg 模式下使用 Keras 的工具
                from tensorflow.keras.utils import load_img, img_to_array  # 再次确保延迟导入
                img = load_img(image_path, target_size=(224, 224))  # 加载图像并调整大小到224x224
                img = img_to_array(img)  # 将图像转换为数组
            elif method == 'flat':
                img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)  # 以灰度模式读取图像
                img = cv2.resize(img, (32, 32))  # 调整图像大小到32x32
            batch_images.append(img)  # 将图像数组添加到批次图像列表
            
            # 尝试根据常见的数据组织方式推断标签：
            # 1) 如果图片位于类子文件夹中（如 .../train/dog/xxx.jpg），优先使用父目录名作为类别
            # 2) 否则，回退到根据文件名解析（原有逻辑）
            parent = os.path.basename(os.path.dirname(image_path))
            train_base = os.path.basename(os.path.normpath(train_folder))
            if parent and parent != train_base:
                label_name = parent
            else:
                label_name = os.path.basename(image_path).split('.')[0]

            # 暂时把 label_name 保存在 batch_labels，再在外面统一映射为整数
            batch_labels.append(label_name)

        batch_images = np.array(batch_images)  # 将批次图像列表转换为numpy数组
        if method == 'vgg':
            from tensorflow.keras.applications.vgg16 import preprocess_input  # 延迟导入
            batch_images = preprocess_input(batch_images)  # 对批次图像进行预处理
            batch_pixels = model.predict(batch_images, verbose=0)  # 使用VGG16模型进行预测
        else:
            # 将 batch_images (N, H, W) 展平为 (N, H*W)
            batch_pixels = batch_images.reshape((batch_images.shape[0], -1))

        # 将处理后的图像特征和标签追加到总列表（确保在循环内）
        X.extend(batch_pixels)
        y.extend(batch_labels)

    # 转换为 numpy 数组
    # 将X和y分别转换为 numpy 数组
    X = np.array(X)
    y = np.array(y)

    # 如果 y 中存储的是类别名称字符串（而非整数），映射到整数标签
    if y.size > 0 and (y.dtype.type is np.str_ or y.dtype == object):
        unique_classes = sorted(list({str(v) for v in y}))
        class_to_idx = {c: i for i, c in enumerate(unique_classes)}
        logging.info(f"检测到类别名称，类别映射: {class_to_idx}")
        y = np.array([class_to_idx[str(v)] for v in y], dtype=np.int64)
    
    logging.info(f"X.shape: {X.shape}")  # 打印X的形状
    logging.info(f"y.shape: {y.shape}")  # 打印y的形状

    # 仅当 X 和 y 非空时才写入缓存
    if X.size > 0 and y.size > 0:
        try:
            with open(x_file_path, 'wb') as fx:
                pickle.dump(X, fx)
            with open(y_file_path, 'wb') as fy:
                pickle.dump(y, fy)
        except Exception as e:
            logging.exception(f"写入缓存文件失败: {e}")
    else:
        logging.warning("X 或 y 为空，跳过写入缓存文件。")

    return X, y  # 返回构建的X和y
    finally:
        # 清理临时解压目录（如果有）
        try:
            if extracted_tmp and os.path.exists(extracted_tmp):
                shutil.rmtree(extracted_tmp)
                logging.info(f"清理临时解压目录: {extracted_tmp}")
        except Exception:
            pass

